#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

import json
import click
import os

from Bio import SeqIO
from tqdm import tqdm


@click.command()
@click.argument("json_path", type=click.Path(exists=True))
@click.argument("fasta_path", type=click.Path(exists=True))
@click.argument("highlight_id")
@click.argument("out_path")
def format_mauve_json(json_path, fasta_path, highlight_id, out_path):
    with open(json_path, 'r') as load_f:
        load_dict = json.load(load_f)
    outjson = []
    seq_id_desc = get_fasta_def(fasta_path)
    for sets in load_dict:
        outset = []
        for item in sets:
            item.pop("gaps")
            fasta_id = item["name"].split(".")[0]
            if fasta_id == highlight_id:
                item["hightlight"] = True
            item["name"] = seq_id_desc[fasta_id]
            outset.append(item)
        outjson.append(outset)
    with open(out_path, 'w') as f:
        json.dump(outjson, f, indent=4)

def get_fasta_def(fasta_path: str):
    ht = {}
    for dirs in os.listdir(fasta_path):
        dir_path = os.path.join(fasta_path, dirs)
        if os.path.isdir(dir_path):
            for item in os.listdir(dir_path):
                file_path = os.path.join(dir_path, item)
                if os.path.isfile(file_path) and file_path.endswith(".fasta"):
                    for seq_record in SeqIO.parse(file_path, "fasta"):
                        seq_id = seq_record.id
                        seq_desc = seq_record.description
                        ht[seq_id] = seq_desc
    return ht

if __name__ == '__main__':
    format_mauve_json()
